import json
from itertools import chain

from ansi2html import Ansi2HTMLConverter
from loguru import logger
from nicegui import ui, app
from tasks.server_helper import ServerHelper
from tasks.model.task import TaskStatus  # Import TaskStatus from our model

conv = Ansi2HTMLConverter(
    line_wrap=True,
    inline=False,
    dark_bg=True,
    escaped=False,
)

field_renderers = {}
live_updates_session = {}


def render_graph_for_task(task_hash: str):
    """
    :param task_hash:
    :return:
    """
    """
    digraph UML_Class_diagram {
    node [
        shape=record
        style=filled
        fillcolor=gray95
    ]
     System_1 [
        shape=plain
        label=<<table border="1" cellborder="0" cellspacing="0" cellpadding="1">
            <tr> <td> <b>System</b> </td> </tr>
            <tr> <td align="left" >+ property</td> </tr>
            <tr> <td port="ss1" align="left" >- p1</td> </tr>
            <tr> <td port="ss2" align="left" >- p2</td> </tr>
            <tr> <td port="ss3" align="left" >- p3</td> </tr>
            <tr> <td align="left">...</td> </tr>
        </table>>
    ]

    System_1:ss2 -> Subsystem_2;
    System_1:ss3 -> Subsystem_3;

    }
    """
    pass


def register_renderer(field_name: str):
    def wrapper(cls):
        field_renderers[field_name] = cls
        return cls

    return wrapper


@register_renderer("text/html")
def render_text_html(text: str):
    return ui.html(text)


@register_renderer("image/png")
def render_png(text: str):
    # with ui.element(tag="div"):
    ui.image("data:image/png;base64," + text).classes("min-width-96 max-w-xl")


@register_renderer("text/plain")
def render_plain(text: str):
    with ui.element(tag="pre") as ele:
        ui.html(text)
        return ele


@register_renderer("text/json")
def render_json(text: str):
    with ui.json_editor({"content": {"json": json.loads(text)}}) as ele:
        return ele


@register_renderer("text/mermaid")
def render_mermaid(text: str):
    # ui.code(text, language="mermaid").classes("my-5")
    return ui.mermaid(text)


@register_renderer("text/markdown")
def render_markdown(text: str):
    return ui.markdown(text)


@register_renderer("stdout")
def render_stdout(text: str):
    return ui.html(conv.convert(text.strip()))


@register_renderer("stderr")
def render_stderr(text: str):
    return None
    return ui.html(conv.convert(text.strip()))


def ui_task_detail(task_id):
    assert task_id is not None
    assert isinstance(task_id, str)
    logger.info(f"ui_task_detail {task_id}")
    _ui_task_detail(task_id)


def render_single_element(item):
    if item["output_type"] == "display_data":
        for mime_type, v in item["data"].items():
            # logger.info(f"mime_type: {mime_type}")
            field_renderers[mime_type](v)
    elif item["output_type"] == "stream":
        render_stdout(item["text"])
    elif item["output_type"] == "error":
        render_stdout("\n".join(item["traceback"]))


def get_task_logs(server_helper, task_id, limit=20, offset=0):
    """Get task logs by looking at the task hash and parsing log files."""
    log_files = server_helper.list_files(task_id, "log")
    log_files.sort()
    log_files = log_files[offset : offset + limit]
    logs = []

    if log_files:
        # Create a temporary directory for log files
        import tempfile
        import os
        import json

        with tempfile.TemporaryDirectory() as temp_dir:
            for log_file in log_files:
                log_file_name = log_file.split("/")[-1]
                local_path = os.path.join(temp_dir, log_file_name)

                try:
                    server_helper.download_file(log_file, local_path)
                    with open(local_path, "r") as f:
                        log_data = json.load(f)
                        logs.append(log_data)
                except Exception as e:
                    logger.exception(f"Error reading log file {log_file}: {str(e)}")

    # Apply offset and limit
    return logs


@ui.refreshable
def _ui_task_detail(task_id, element=None):
    # Get server_helper from app context
    server_helper: ServerHelper = app.server_helper

    # Get task details
    task_spec = server_helper.get_task_spec(task_id)

    # Show task information
    # Add CSS to remove minimum height
    ui.add_head_html(
        """
        <style>
            .ag-center-cols-viewport {
                min-height: unset !important;
            }
        </style>
    """
    )

    # Set status color based on task status
    status_color = "text-blue-500"
    if hasattr(task_spec, "status"):
        if task_spec.status == TaskStatus.COMPLETED:
            status_color = "text-green-500"
        elif task_spec.status == TaskStatus.FAILED:
            status_color = "text-red-500"
        elif task_spec.status == TaskStatus.RUNNING:
            status_color = "text-blue-500"

    # Status with color applied
    status_value = task_spec.status.value if hasattr(task_spec, "status") else "N/A"
    status_html = f'<span class="{status_color}">{status_value}</span>'

    # Prepare row data for AgGrid
    info_row_data = [
        {"property": "Task Hash", "value": task_id},
        {
            "property": "Function",
            "value": task_spec.task if hasattr(task_spec, "task") else "N/A",
        },
        {"property": "Status", "value": status_html},
    ]
    with ui.row().classes("w-full"):
        ui.aggrid(
            {
                "columnDefs": [
                    {
                        "headerName": "Property",
                        "field": "property",
                        "width": 120,
                    },
                    {
                        "headerName": "Value",
                        "field": "value",
                    },
                ],
                "rowData": info_row_data,
                "domLayout": "autoHeight",
            },
            html_columns=[1],  # Make the 'value' column render HTML
        ).classes("dense w-full").style("height: 100%")

        with ui.expansion("Task Files", icon="folder").classes("w-full dense"):
            # Get both input files and output files
            input_files = server_helper.list_files(task_id, "file")
            output_files = server_helper.list_files(task_id, "output")

            all_files = []
            for file_path in input_files:
                all_files.append(
                    {
                        "file_id": file_path,
                        "type": "Input File",
                        "name": file_path.split("/")[-1],
                    }
                )

            for file_path in output_files:
                all_files.append(
                    {
                        "file_id": file_path,
                        "type": "Output File",
                        "name": file_path.split("/")[-1],
                    }
                )

            if not all_files:
                ui.label("No files found")
            else:

                def on_click(event):
                    file_id = event.args["data"]["file_id"]
                    column_id = event.args["colId"]
                    if column_id in ["file_id", "name", "download"]:
                        ui.navigate.to(f"/file/{file_id}")
                    elif column_id == "copy_link":
                        link = f"/file/{file_id}"
                        ui.clipboard.write(link)
                        ui.notify(f"File viewer link copied to clipboard")

                row_data = all_files
                for item in row_data:
                    item["download"] = "view"
                    item["copy_link"] = "link"

                ui.aggrid(
                    {
                        "columnDefs": [
                            {
                                "headerName": "File Name",
                                "field": "name",
                                "sortable": True,
                            },
                            {"headerName": "Type", "field": "type", "sortable": True},
                            {
                                "headerName": "File ID",
                                "field": "file_id",
                                "sortable": True,
                            },
                            {
                                "headerName": "View",
                                "field": "download",
                                "sortable": False,
                                "width": 30,
                            },
                            {
                                "headerName": "Copy Link",
                                "field": "copy_link",
                                "sortable": False,
                                "width": 30,
                            },
                        ],
                        "rowData": row_data,
                        "domLayout": "autoHeight",
                    },
                    html_columns=[3, 4],  # Make view and copy link columns clickable
                ).classes("dense").on("cellClicked", on_click).style("height: 100%")

    # Output/logs section
    with ui.card().tight().classes("w-full"):
        ui.label("Task Logs").classes(" font-bold")

        call_dict = dict(
            offset=0,
            limit=20,
        )
        task_ids = []
        if task_spec.task == "__collect__":
            for k, v in task_spec.output_params.items():
                task_ids.append((k, v))
        else:
            task_ids.append((None, task_id))

        def runrun():
            try:
                show_name, task_id = task_ids[0]
                if call_dict["offset"] == 0 and show_name is not None:
                    ui.label(f"{show_name}").classes("font-bold")
                    ui.separator()
                task_logs = get_task_logs(server_helper, task_id, **call_dict)
                cnt_this_time = 0
                for record in task_logs:
                    render_single_element(record)
                    cnt_this_time += 1
                call_dict["offset"] += cnt_this_time

                # Check if task is still running to continue updating logs
                task_status = None
                try:
                    updated_task = server_helper.get_task_spec(task_id)
                    if hasattr(updated_task, "status"):
                        task_status = updated_task.status
                except Exception as e:
                    logger.exception(f"Error getting task status: {str(e)}")
                    pass

                if cnt_this_time == 0 and (
                    task_status is None or task_status != TaskStatus.RUNNING
                ):
                    logger.info("task finished")
                    call_dict["offset"] = 0
                    task_ids.pop(0)
                    if len(task_ids) == 0:
                        timer.cancel()
            except Exception as e:
                logger.exception(f"Error getting logs: {str(e)}")
                timer.cancel()

        timer = ui.timer(0.5, runrun)


page_config = {
    "nav_entry": [
        {
            "title": None,
            "url": "/task_detail/{task_id}",
            "function": ui_task_detail,
        }
    ]
}
